import numpy as np
import torch
import torch.nn.functional as F

from attack.base_attack import BaseAttack

class Attack(BaseAttack):
    def __init__(self, model, config):
        super().__init__('BLB', model, config, batch_size = 100, targeted = True, llc = False)

    def attack_achieved(self, pre_softmax, target_class):
        return np.argmax(pre_softmax) == target_class

    def attack_batch(self, images, targets):
        print("\n") 

        assert len(images) == self.batch_size  

        """
        强制使得两者大小相等,因此见base_method中batch_size的设定
        因为（len(images)=1000）,所以需要执行10遍（batch_size=100）
        未来可在不同攻击下进行参数的更改！
        """

        var_images = torch.from_numpy(images).to(self.model.device)
        var_targets = torch.LongTensor(targets).to(self.model.device)

        #定义参数范围便于搜索
        const_origin = np.ones(shape = self.batch_size, dtype = float) * self.config['init_const']
        const_upper_bound = [1e10] * self.batch_size
        const_lower_bound = np.zeros(self.batch_size)

        # L2攻击（具体含义需要再看）
        best_prediction_class = [1e10] * self.batch_size
        adv_images = np.zeros(var_images.size())
        current_prediction_class = [-1] * self.batch_size

        self.model.eval()

        output = self.model(var_images)
        loss = F.cross_entropy(output,var_targets,reduction = 'none') 
        #对于batch的images应该关闭reduction，https://blog.csdn.net/goodxin_ie/article/details/89645358
        #分别对每个图像进行loss函数的返回，不作为整体使用

        for search_for_const in range(self.config['second_layer_loop']):
            per = torch.zeros_like(var_images).float().cpu().numpy() #产生扰动(tensor->numpy)
            per = torch.from_numpy(per).to(self.model.device)
            per.requires_grad = True

            # 用LBFGS去操作per扰动，用默认的的参数
            optimizer = torch.optim.LBFGS([per], max_iter = self.config['max_iterations'])
            var_const = torch.from_numpy(const_origin).to(self.model.device) #不可移出循环
            print("\tsecond_layer_loop {}:".format(search_for_const))

            def loss():       #计算损失函数流程（参考论文）
                per_images = torch.clamp(var_images + per, min=0.0, max=1.0)
                predict = self.model(per_images)
                l2dist = torch.sum((per_images - var_images) ** 2, [1, 2, 3])
                pre_loss = F.cross_entropy(predict, var_targets)
                pre_pro_loss = var_const * pre_loss
                loss = l2dist.sum() + pre_pro_loss.sum()  # min const*|pre| + pre_pro_loss(x+pre,label)
                optimizer.zero_grad()
                loss.backward(retain_graph=True)
                return loss

            optimizer.step(loss) #每次改变一次optimize的参数
        
            #参数提出
            per_images = torch.clamp(var_images + per, min=0.0, max=1.0)
            predict = self.model(per_images)
            l2dist = torch.sum((per_images - var_images) ** 2, [1, 2, 3])

            #此后过程类似CW2攻击（联系）
            for i, (dist, score, perturbation) in enumerate(
                    zip(l2dist.data.cpu().numpy(), predict.data.cpu().numpy(), per_images.data.cpu().numpy())):
                if dist < best_prediction_class[i] and self.attack_achieved(score, targets[i]):
                    best_prediction_class[i] = dist
                    current_prediction_class[i] = np.argmax(score)
                    adv_images[i] = perturbation

            # 每次修改常量const_origin
            for i in range(self.batch_size):
                if current_prediction_class[i] == targets[i] and current_prediction_class[i] != -1:
                    const_upper_bound[i] = min(const_upper_bound[i], const_origin[i])
                    if const_upper_bound[i] < 1e10:
                        const_origin[i] = (const_lower_bound[i] + const_upper_bound[i]) / 2.
                else:
                    const_lower_bound[i] = max(const_lower_bound[i], const_origin[i])
                    if const_upper_bound[i] < 1e10:
                        const_origin = (const_lower_bound[i] + const_upper_bound[i]) / 2
                    else:
                        const_origin[i] *= 10
    
        return np.array(adv_images).astype(np.float32)